"""Sparse encodes the k-fold cross validation training data for given protein
groups.

Creates k sparse encoded training files and k sparse encoded test files in a
k-fold cross validation fashion. The variable k depends on the number of groups
in the given group file. All the files are outputted to given outputdir. The sparse encoding depends on the options that are supplied to the program. There are no default options.
Options:
    -s      Position-Specific Score Matrix
    -w      Weighted Observed Percentages Rounded Down 
    -i      Information Per Position
    -g      Relative Weight of Gapless Real Matches to Pseudocounts
    -x int  Windowx where window size = 2 * x + 1
    -a      Local sequence information
Usage:
    python sparseencode_trainingdata.py -six 5 ../results/groupoutput ../results/testoutput
"""
import sys
import getopt
import random
import math

class TrainingPoint:
    """A training point. Its attributes are the target value, the local
    sequence information and the profile features. A TrainingPoint can return a
    sparse encoded version of itself."""
    def __init__(self):
        self.aafeatures = None
        self.target = None
        self.profilefeatures = None

    def set_aa_window_size(self, windowsize):
        """"Initializes the local sequence information features."""
        self.aafeatures = [0] * 20 * windowsize

    def set_profile_window_size(self, numfeatures, windowsize):
        """"Initializes the profile information features."""
        self.profilefeatures = [0] * numfeatures * windowsize

    def get_aa_index(self, aa):
        """Returns the index of a given amino acid."""
        if aa == 'A': return 0
        elif aa == 'R': return 1
        elif aa == 'N': return 2
        elif aa == 'D': return 3
        elif aa == 'C': return 4
        elif aa == 'E': return 5
        elif aa == 'Q': return 6
        elif aa == 'G': return 7
        elif aa == 'H': return 8
        elif aa == 'I': return 9
        elif aa == 'L': return 10
        elif aa == 'K': return 11
        elif aa == 'M': return 12
        elif aa == 'F': return 13
        elif aa == 'P': return 14
        elif aa == 'S': return 15
        elif aa == 'T': return 16
        elif aa == 'W': return 17
        elif aa == 'Y': return 18
        elif aa == 'V': return 19
        else: sys.stderr.write('Undefined AA')

    def set_aa_in_window(self, aa, windowindex):
        """Indicate that the amino acid aa is at location windowindex"""
        self.aafeatures[self.get_aa_index(aa) + windowindex * 20] = 1

    def set_profile_in_window(self, profile, windowindex):
        """Indicate that the amino acid aa at location windowindex has profile
        profile"""
        for i in range(len(profile)):
            self.profilefeatures[i + windowindex * len(profile)] = profile[i]
            
    def sparse_encode(self):
        """Return a sparse encoded string of this training point."""
        rv = self.target

        if self.aafeatures != None:
            for i in range(len(self.aafeatures)):
                rv += ' ' + str(i + 1) + ':' + str(self.aafeatures[i])

        if self.profilefeatures != None:
            for i in range(len(self.profilefeatures)):
                if self.aafeatures != None:
                    rv += ' ' + str(len(self.aafeatures) + i + 1) + ':' + \
                            str(self.profilefeatures[i])
                else:
                    rv += ' ' + str(i + 1) + ':' + \
                            str(self.profilefeatures[i])

        return rv

    def set_target_value(self, value):
        self.target = value

def sparse_encode(z_coord, aa_window=None, profile_window=None):
    """Sparse encodes a window of amino acids and/or profiles"""
    tp = TrainingPoint()

    if aa_window != None:
        tp.set_aa_window_size(len(aa_window))

        for i in range(len(aa_window)):
            if aa_window[i] != 'X':
                tp.set_aa_in_window(aa_window[i], i)

    if profile_window != None:
        tp.set_profile_window_size(len(profile_window[0]), len(profile_window))

        for i in range(len(profile_window)):
            tp.set_profile_in_window(profile_window[i], i)

    z_coord = float(z_coord)
    if abs(z_coord) <= 22 and abs(z_coord) > 10:
        tp.set_target_value('+1')
    else:
        tp.set_target_value('-1')

    return tp.sparse_encode()


def sparse_encode_protein(protein, windowx, options):
    """Reads the sequence with name protein from the membrane_training_data
    folder and the profile with name protein from the profiles folder. The
    sparse encoding depends on the given options."""
    rv = ''

    pfh = open('../data/membrane_training_data/' + protein + '.txt')
    aas = pfh.readlines()[1:]

    # Store profilefeatures if specified
    profilefeatures = []
    if "pss" in options or "wp" in options or "pi" in options or "gw" in options:
        profilesfile = open('../data/profiles/' + protein + \
                '.psi').readlines()[3:-6]
        for line in profilesfile:
            columns = line.split()

            selectedcolumns = []
            if "pss" in options:
                normalizedscores = map(lambda x: 1 / (1 + math.pow(math.e, -int(x))), columns[2:22])
                selectedcolumns.extend(normalizedscores)
            if "wp" in options:
                selectedcolumns.extend(columns[22:42])
            if "pi" in options:
                selectedcolumns.append(columns[42])
            if "gw" in options:
                selectedcolumns.append(columns[43])
            profilefeatures.append(tuple(selectedcolumns))

    # Create windows of local sequence information and profile information
    # depending on specified options.
    for i in range(0, len(aas)):
        aa_window = []
        profile_window = []

        for j in range(-windowx, windowx + 1):
            if i + j >= 0 and i + j < len(aas):
                columns = aas[i + j].split()
                aa_window.append(columns[1])
                if len(profilefeatures) > 0:
                    profile_window.append(profilefeatures[i + j])
            else:
                aa_window.append('X')
                if len(profilefeatures) > 0:
                    profile_window.append(tuple([0] * len(profilefeatures[0])))

        if "pss" in options or "wp" in options or "pi" in options or "gw" in options:
            if "aa" in options:
                rv += sparse_encode(aas[i].split()[6], profile_window=profile_window, aa_window=aa_window) + '\n'
            else:
                rv += sparse_encode(aas[i].split()[6], profile_window=profile_window) + '\n'
        elif "aa" in options:
            rv += sparse_encode(aas[i].split()[6], aa_window=aa_window) + '\n'

    pfh.close()

    return rv

def process(groupoutput, outputdir, windowx, options):
    """Reads the groups from groupoutput and creates trainingdata and testdata
    in outputdir."""
    gofh = open(groupoutput, 'r')

    groupnumber = -1
    groups = []

    for line in gofh:
        if line[:5] == 'Group':
            groupnumber = int(line[5:])
            groups.append([])
        else:
            groups[groupnumber].append(line[:-1])
    gofh.close()

    se_groups = []
    for g in groups:
        se_groups.append(map(lambda p: sparse_encode_protein(p, windowx, options), g))

    for testi in range(len(se_groups)):
        trainfh = open(outputdir + 'traink' + 
                str(len(se_groups)) + 't' + str(testi) + '.dat', 'w')
        testfh = open(outputdir + 'testk' + 
                str(len(se_groups)) + 't' + str(testi) + '.dat', 'w')

        for j in range(len(se_groups)):
            if testi != j:
                for p in se_groups[j]:
                    trainfh.write(p)
            else:
                for p in se_groups[j]:
                    testfh.write(p)

        trainfh.close()     
        testfh.close()

def main():
    # parse command line options
    try:
        opts, args = getopt.getopt(sys.argv[1:], "ahswigx:", ["help"])
    except getopt.error, msg:
        print msg
        print "for help use --help"
        sys.exit(2)
    # process options
    options = []
    windowx = None
    for o, a in opts:
        print o
        if o in ("-h", "--help"):
            print __doc__
            sys.exit(0)
        elif o == "-s":
            options.append('pss')
        elif o == "-w":
            options.append('wp')
        elif o == "-i":
            options.append('pi')
        elif o == "-g":
            options.append('gw')
        elif o == "-x":
            windowx = int(a)
        elif o == "-a":
            options.append('aa')
    print options
    print windowx
    # process arguments
    if windowx == None or len(options) == 0 or len(args) != 2:
        print __doc__
    else:
        process(args[0], args[1], windowx, options) # process() is defined elsewhere

if __name__ == "__main__":
    main()
